import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'
import { Document } from 'mongoose'

export interface TemplateMetadata {
    type: 'sdxl' | 'controlnet' | 'lora'
    controlType?: string
    loraName?: string
    loraWeight?: number
    tags: string[]
    author: string
    version: string
    createTime: Date
    updateTime: Date
}

export interface TemplateParams {
    prompt: string
    negativePrompt: string
    width: number
    height: number
    seed?: number
    sampler: {
        name: string
        scheduler?: string
        steps: number
        cfg: number
        denoise?: number
        seed?: number
    }
    controlnet?: {
        type: string
        strength: number
        image?: string
        preprocessor?: {
            low_threshold?: number
            high_threshold?: number
        }
    }
    lora?: {
        name: string
        strength: number
    }
}

@Schema({ timestamps: true })
export class WorkflowTemplate {
    @Prop({ required: true, unique: true })
    id: string

    @Prop({ required: true })
    name: string

    @Prop()
    description?: string

    @Prop({ required: true })
    category: string

    @Prop({ type: Object, required: true })
    workflow: any

    @Prop({ type: Object, required: true })
    defaultParams: TemplateParams

    @Prop({ type: Object, required: true })
    metadata: TemplateMetadata
}

export type WorkflowTemplateDocument = WorkflowTemplate & Document
export const WorkflowTemplateSchema = SchemaFactory.createForClass(WorkflowTemplate)